# -*- coding: utf-8 -*-
from scipy import *
from scipy import interpolate

'''
module to construct the interpolation function from external data
The interpolation data contains the values for
theta 0 pi/2
phi 0 pi
The function is extended to larger range using the symmetry of the function.
'''

def stable_check(A, theta):
    # Check for bistable state for SW particle
    # A hack to revert the theta value into the 0 to pi/2 range.
    theta = abs(pi*floor(theta/(pi/2))-theta)
    t = tan(theta)**(1./3.)
    t2=t*t
    t4=t2*t2
    front = sqrt(1+t2+t4)/(1+t2)
    return A<front

int_filename='SW.dat'
# we start to load the data directly
int_data = loadtxt(int_filename)
alist = int_data[:,0]
tlist = int_data[:,1]
p1list = int_data[:,2]
p2list = int_data[:,3]
e1list = int_data[:,4]
e2list = int_data[:,5]
e3list = int_data[:,6]

#test_alist = arange(0,0.99,0.04)
#test_tlist = arange(0,pi/2+0.00001,pi/40)
#test_p1list = p1list.copy()
#test_p1list.resize((len(test_alist),len(test_tlist)))
#test_p2list = p2list.copy()
#test_p2list.resize((len(test_alist),len(test_tlist)))

coord = stable_check(alist,tlist)
select_alist=alist[coord]
select_tlist=tlist[coord]
select_p2list=p2list[coord]
select_e2list=e2list[coord]
select_e3list=e3list[coord]


#p1=interpolate.RectBivariateSpline(test_alist,test_tlist,test_p1list)
#p2=interpolate.RectBivariateSpline(test_alist,test_tlist,test_p2list)
p1=interpolate.interp2d(alist,tlist,p1list,kind='cubic')
p2=interpolate.interp2d(select_alist,select_tlist,select_p2list,kind='cubic')
e1=interpolate.interp2d(alist,tlist,e1list,kind='cubic')
e2=interpolate.interp2d(select_alist,select_tlist,select_e2list,kind='cubic')
e3=interpolate.interp2d(select_alist,select_tlist,select_e3list,kind='cubic')

# extend the interpolation function using symmetry
def np1(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = 2*pi - p1(a, theta)[0]
    else:
        out = p1(a, theta)[0]
    return out

def np2(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = 2*pi - p2(a, theta)[0]
    else:
        out = p2(a, theta)[0]
    return out

def ne1(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = e1(a, theta)[0]
    else:
        out = e1(a, theta)[0]
    return out

def ne2(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = e2(a, theta)[0]
    else:
        out = e2(a, theta)[0]
    return out

def ne3(a, theta):
    if theta > pi/2:
        theta = pi - theta
        out = e3(a, theta)[0]
    else:
        out = e3(a, theta)[0]
    return out

if __name__=='__main__':
    print ne2(0.2,0.3)